﻿namespace Cellfish.Toolbox.Tasks
{
    using System;
    using System.Collections.Generic;
    using System.Linq;
    using System.Threading.Tasks;

    /// <summary>
    /// Provides task combinators.
    /// </summary>
    public static partial class TaskWhen
    {
        /// <summary>
        /// Creates a task that will complete when all of the supplied tasks have completed or when the first task fails.
        /// </summary>
        /// <param name="tasks">The tasks to wait on for completion.</param>
        /// <returns>A task that represents the completion of all of the supplied tasks.</returns>
        /// <exception cref="System.ArgumentNullException">If task is null.</exception>
        public static Task AllOrError(params Task[] tasks)
        {
            return AllOrError((tasks ?? new Task[0]).AsEnumerable());
        }

        /// <summary>
        /// Creates a task that will complete when all of the supplied tasks have completed or when the first task fails.
        /// </summary>
        /// <typeparam name="T">The type of the completed task.</typeparam>
        /// <param name="tasks">The tasks to wait on for completion.</param>
        /// <returns>A task that represents the completion of all of the supplied tasks.</returns>
        /// <exception cref="System.ArgumentNullException">If task is null.</exception>
        public static Task<T[]> AllOrError<T>(params Task<T>[] tasks)
        {
            return AllOrError((tasks ?? new Task<T>[0]).AsEnumerable());
        }

        /// <summary>
        /// Creates a task that will complete when all of the supplied tasks have completed or when the first task fails.
        /// </summary>
        /// <param name="tasks">The tasks to wait on for completion.</param>
        /// <returns>A task that represents the completion of all of the supplied tasks.</returns>
        /// <exception cref="System.ArgumentNullException">If task is null.</exception>
        public static Task AllOrError(IEnumerable<Task> tasks)
        {
            if (tasks == null)
            {
                throw new ArgumentNullException("tasks");
            }

            var errorResult = new TaskCompletionSource<bool>();
            var tasksToCheck = tasks.Select(
                task => task.ContinueWith(
                    t =>
                        {
                            if (t.IsCanceled)
                            {
                                errorResult.TrySetCanceled();
                            }
                            else if (t.IsFaulted)
                            {
                                errorResult.TrySetException(t.Exception.InnerException);
                            }
                        }));

            return Task.WhenAny(errorResult.Task, Task.WhenAll(tasksToCheck)).Unwrap();
        }

        /// <summary>
        /// Creates a task that will complete when all of the supplied tasks have completed or when the first task fails.
        /// </summary>
        /// <typeparam name="T">The type of the completed task.</typeparam>
        /// <param name="tasks">The tasks to wait on for completion.</param>
        /// <returns>A task that represents the completion of all of the supplied tasks.</returns>
        /// <exception cref="System.ArgumentNullException">If task is null.</exception>
        public static Task<T[]> AllOrError<T>(IEnumerable<Task<T>> tasks)
        {
            if (tasks == null)
            {
                throw new ArgumentNullException("tasks");
            }

            var errorResult = new TaskCompletionSource<T[]>();
            var tasksToCheck = tasks.Select(
                task => task.ContinueWith(
                    t =>
                    {
                        if (t.IsCanceled)
                        {
                            errorResult.TrySetCanceled();
                        }
                        else if (t.IsFaulted)
                        {
                            errorResult.TrySetException(t.Exception.InnerException);
                        }
                        else
                        {
                            return t.Result;
                        }

                        return default(T);
                    }));

            return Task.WhenAny(errorResult.Task, Task.WhenAll(tasksToCheck)).Unwrap();
        }
    }
}
